#include "tree.h"
#include "alignment.h"
#include "MbRandom.h"
#include <iostream>
#include <iomanip>
#include <string>
#include <vector>
#include <istream>
#include <sstream>
#include <fstream>
#include <stdio.h>
#include <omp.h>


void   buildTreeFromNewickString(Tree * thisTree, string treeStr);
void   buildRandomTree(Tree * thisTree);
int   dex(Node *p);
void   passDn(Tree * thisTree, Node *p, int *x);
void   showNodes(Tree * thisTree, Node *p, int indent);

Node *  createNode() {
    Node * thisNode = (Node *) malloc(sizeof(Node));

	thisNode->lft    = NULL;
	thisNode->rht    = NULL;
	thisNode->anc    = NULL;
	thisNode->index  = 0;
	thisNode->isLeaf = false;
	thisNode->flag   = false;
	thisNode->name   = (char *) malloc(sizeof(char)*STRLEN);

    *(thisNode->name) = '\0';
    thisNode->childrenLeft = 0;
    return thisNode;
}

Node *  createNodes(int numNodes) {
    Node * thisNode = (Node *) malloc(numNodes*sizeof(Node));

    int i;
    for(i=0; i<numNodes; i++)
    {
        thisNode[i].lft    = NULL;
        thisNode[i].rht    = NULL;
        thisNode[i].anc    = NULL;
        thisNode[i].index  = 0;
        thisNode[i].isLeaf = false;
        thisNode[i].flag   = false;
        thisNode[i].name   = (char *) malloc(sizeof(char)*STRLEN);

            
        thisNode[i].childrenLeft = 0;
    }
    return thisNode;
}

Node **  createNodesPP(int numNodes) {
    Node ** thisNode = (Node **) malloc(numNodes*sizeof(Node *));
    return thisNode;
}



Node   *getLft(Node * thisNode) { return thisNode->lft; }
Node   *getRht(Node * thisNode) { return thisNode->rht; }
Node   *getAnc(Node * thisNode) { return thisNode->anc; }
int   getIndex(Node * thisNode) { return thisNode->index; }
bool   getIsLeaf(Node * thisNode) { return thisNode->isLeaf; }
string   getName(Node * thisNode) { return thisNode->name; }
bool   getFlag(Node * thisNode) { return thisNode->flag; }
double   getV(Node * thisNode) { return thisNode->v; }
void   setLft(Node * thisNode, Node *p) { thisNode->lft = p; }
void   setRht(Node * thisNode, Node *p) { thisNode->rht = p; }
void   setAnc(Node * thisNode, Node *p) { thisNode->anc = p; }
void   setIndex(Node * thisNode, int x) { thisNode->index = x; }
void   setIsLeaf(Node * thisNode, bool x) { thisNode->isLeaf = x; }
void   setName(Node * thisNode, const char * s) { strncpy(thisNode->name, s, STRLEN); }
void   setFlag(Node * thisNode, bool x) { thisNode->flag = x; }
void   setV(Node * thisNode, double x) { thisNode->v = x; }
void   setChildrenLeft(Node * thisNode, int x) { thisNode->childrenLeft = x; }


Tree * createTree(Alignment *al, MbRandom *rp, string treeStr) {
    Tree * thisTree = (Tree *) malloc(sizeof(Tree));

	thisTree->nodes = NULL;
	thisTree->downPassSequence = NULL;
	thisTree->alignmentPtr = al;
	thisTree->ranPtr = rp;
    thisTree->isRooted = false;
	
	if (treeStr == "")
		buildRandomTree(thisTree);
	else
		buildTreeFromNewickString(thisTree, treeStr);
    return thisTree;
	
}

void deleteTree(Tree * thisTree) {
    
	if (thisTree->nodes != NULL)
		free(thisTree->nodes);
	if (thisTree->downPassSequence != NULL)
		free(thisTree->downPassSequence);
    
}

int   getNumTaxa(Tree * thisTree) { return thisTree->numTaxa; }
int   getNumNodes(Tree * thisTree) { return thisTree->numNodes; }

void buildRandomTree(Tree * thisTree) {
    
	thisTree->numTaxa = thisTree->alignmentPtr->getNumTaxa();
    //fprintf(stderr,"BEFORE: num taxa %d\n", thisTree->numTaxa);
	if (thisTree->isRooted) {
        //fprintf(stderr,"isRooted: numtaxa: %d\n", thisTree->numTaxa);
        //fprintf(stderr,"RHS: %d\n", 2 * thisTree->numTaxa);
		thisTree->numNodes = 2 * thisTree->numTaxa;
        //fprintf(stderr,"LHS: %d\n", thisTree->numNodes);
    }
	else {
        //fprintf(stderr,"NOT isRooted: numtaxa: %d\n", thisTree->numTaxa);
		//fprintf(stderr,"RHS: %d\n", 2 * thisTree->numTaxa-2);
        thisTree->numNodes = 2 * thisTree->numTaxa - 2;
        //fprintf(stderr,"LHS: %d\n", thisTree->numNodes);
    }
    //fprintf(stderr,"AFTER: num taxa %d, num nodes %d, calc'ed %d\n", thisTree->numTaxa, thisTree->numNodes, 2 * thisTree->numTaxa - 2);
    //exit(0);
	/* allocate the nodes for the tree */
	thisTree->nodes = createNodes(thisTree->numNodes);
    
    for (int i=0; i<thisTree->numNodes; i++)
		setIndex(&(thisTree->nodes[i]), i);
	thisTree->downPassSequence = createNodesPP(thisTree->numNodes);
    thisTree->nextQueue = createNodesPP(thisTree->numNodes);
    thisTree->queueStart = 0;
    thisTree->queueEnd = 0;
	for (int i=0; i<thisTree->numNodes; i++)
    {
		thisTree->downPassSequence[i] = NULL;
        thisTree->nextQueue[i] = NULL;
    }
	for (int i=0; i<thisTree->numTaxa; i++)
		setName(&(thisTree->nodes[i]), thisTree->alignmentPtr->getTaxonName(i).c_str() );
    
	Node **availableNodes = createNodesPP(thisTree->numNodes);
	int numAvailableNodes = 0;
	int nextTip = 0;
	int nextInt = thisTree->numTaxa;
	
	/* set up three taxon tree */
	int numTaxaInTree = 0;
	Node *p;
	if (thisTree->isRooted == true)
		p = &thisTree->nodes[nextInt++];
	else
    {
		p = &thisTree->nodes[nextTip++];
		numTaxaInTree++;
    }
	thisTree->root = p;
	Node *q = &thisTree->nodes[nextInt++];
	availableNodes[numAvailableNodes++] = q;
	setLft(p, q);
	setAnc(q, p);
	p = q;
	q = &thisTree->nodes[nextTip++];
	numTaxaInTree++;
	availableNodes[numAvailableNodes++] = q;
	setLft(p, q);
	setAnc(q, p);
	q = &thisTree->nodes[nextTip++];
	numTaxaInTree++;
	availableNodes[numAvailableNodes++] = q;
	setRht(p, q);
	setAnc(q, p);
	
	int numTips = nextTip;
	for (int i=numTips; i<thisTree->numTaxa; i++)
    {
		p = availableNodes[(int)(thisTree->ranPtr->uniformRv()*numAvailableNodes)];
		q = getAnc(p);
		Node *a = &thisTree->nodes[nextTip++];
		numTaxaInTree++;
		Node *b = &thisTree->nodes[nextInt++];
		if (getLft(q) == p)
        {
			setLft(q, b);
			setAnc(b, q);
			setLft(b, p);
			setAnc(p, b);
			setRht(b, a);
			setAnc(a, b);
        }
		else
        {
			setRht(q, b);
			setAnc(b, q);
			setRht(b, p);
			setAnc(p, b);
			setLft(b, a);
			setAnc(a, b);
        }
		availableNodes[numAvailableNodes++] = a;
		availableNodes[numAvailableNodes++] = b;
    }
	delete [] availableNodes;
    
	/* initialize the post-order tree traversal sequence */
	getDownPassSequence(thisTree);
	//fprintf(stderr,"here1\n");
	for (int n=0; n<thisTree->numNodes; n++)
    {
		Node *p = getDownPassNode( thisTree, n );
		if (getAnc(p) != NULL)
        {
            //fprintf(stderr,"trying w/anc %d of %d...",n, thisTree->numNodes);
			setV(p, thisTree->ranPtr->exponentialRv(10.0) );
            //fprintf(stderr,"done\n");
        }
		else
        {
            //fprintf(stderr,"trying w/out %d of %d...",n, thisTree->numNodes);
			setV(p, 0.0 );
            //fprintf(stderr,"done\n");
        }
    }
    //fprintf(stderr,"here2\n");
#	if 0
	showNodes(thisTree->root, 3);
	cout << "numTaxaInTree = " << numTaxaInTree << endl;
#	endif
}

/* build a tree up from a string with the Newick description of the tree */
void buildTreeFromNewickString(Tree * thisTree, string treeStr) {
    
	/* place all of the tree tokens into a string vector, and count the number of taxa in the tree while we are at it */
	vector<string> words;
	thisTree->numTaxa = 0;
	int numLeftParantheses = 0;
	string nameStr = "";
	bool readingLength = false;
	bool brlensPresent = false;
	for (int i=0; i<treeStr.length(); i++)
    {
		char ch = treeStr.at(i);
		if ( ch == '(' || ch == ')' || ch == ',' || ch == ':' || ch == ';' )
        {
			string tmp = "";
			tmp += ch;
			words.push_back(tmp);
			nameStr = "";
			if ( ch == ':' )
            {
				readingLength = true;
				brlensPresent = true;
            }
			else if ( ch == '(' )
				numLeftParantheses++;
			//cout << "\"" << ch << "\"" << endl;
        }
		else
        {
			nameStr += ch;
			if ( treeStr.at(i+1) == '(' || treeStr.at(i+1) == ')' || treeStr.at(i+1) == ',' || treeStr.at(i+1) == ':' || treeStr.at(i+1) == ';' )
            {
				words.push_back(nameStr);
				if ( readingLength == false )
					thisTree->numTaxa++;
				else
					readingLength = false;
            }
        }
    }
	/* decide whether the tree is rooted or unrooted from the number of taxa and the number of left parantheses */
	if ( numLeftParantheses == thisTree->numTaxa - 1 )
		thisTree->isRooted = true;
	else if ( numLeftParantheses == thisTree->numTaxa - 2 )
		thisTree->isRooted = false;
	else
    {
		cerr << "Inconsistent number of parantheses in tree description" << endl;
		exit(1);
    }
    
	/* calculate the number of nodes in the tree */
	if (thisTree->isRooted == true)
		thisTree->numNodes = 2 * thisTree->numTaxa;
	else
		thisTree->numNodes = 2 * thisTree->numTaxa - 2;
    
	/* check for consistency in the number of taxa between the data and the tree */
	if ( thisTree->numTaxa != thisTree->alignmentPtr->getNumTaxa() )
    {
		cerr << "Problem reading the correct number of taxa." << endl;
		exit(1);
    }
    
	/* allocate the nodes for the tree */
	thisTree->nodes = createNodes(thisTree->numNodes);
	for (int i=0; i<thisTree->numNodes; i++)
		setIndex(&(thisTree->nodes[i]), i);
	thisTree->downPassSequence = createNodesPP(thisTree->numNodes);
    thisTree->nextQueue = createNodesPP(thisTree->numNodes);
    thisTree->queueStart = 0;
    thisTree->queueEnd = 0;
	for (int i=0; i<thisTree->numNodes; i++)
    {
		thisTree->downPassSequence[i] = NULL;
        thisTree->nextQueue[i] = NULL;
    }
	/* build the tree */
	int nextIntNode = thisTree->numTaxa;
	readingLength = false;
	Node *p, *q;
	for (vector<string>::iterator s=words.begin(); s != words.end(); s++)
    {
		if ( (*s) == "(" )
        {
			/* add an internal node */
			if (nextIntNode == thisTree->numTaxa)
            {
				p = &thisTree->nodes[nextIntNode++];
				thisTree->root = p;
            }
			else
            {
				if (getLft(p) == NULL)
                {
					q = &thisTree->nodes[nextIntNode++];
					setAnc(q, p);
					setLft(p, q);
					p = q;
                }
				else if (getRht(p) == NULL)
                {
					q = &thisTree->nodes[nextIntNode++];
					setAnc(q, p);
					setRht(p, q);
					p = q;
                }
				else if (getAnc(p) == NULL)
                {
					q = &thisTree->nodes[nextIntNode++];
					setLft(q, p);
					setAnc(p, q);
					setFlag(q, true );
					p = q;
					thisTree->root = p;
                }
				else
                {
					cout << "Problem finding an open node" << endl;
					exit(1);
                }
            }
        }
		else if ( (*s) == ")" )
        {
			/* go down */
			if (getFlag(p) == false)
				p = getAnc(p);
			else
				p = getLft(p);
        }
		else if ( (*s) == "," )
        {
			/* go down */
			if (getFlag(p) == false)
				p = getAnc(p);
			else
				p = getLft(p);
        }
		else if ( (*s) == ":" )
        {
			/* start reading a digit */
			readingLength = true;
        }
		else if ( (*s) == ";" )
        {
			/* end of the tree */
            
        }
		else
        {
			if ( readingLength == true )
            {
				/* add a branch length */
				double x;
				istringstream buf( (*s) );
				buf >> x;
				//p->setV( x );
				readingLength = false;
            }
			else
            {
				/* add a tip node */
				int taxonIndex = thisTree->alignmentPtr->getTaxonIndex( (*s) );
				if (taxonIndex == -1)
                {
					cerr << "Couldn't find \"" << (*s) << "\" in the list of taxon names." << endl;
					exit(1);
                }
				
				if (getLft(p) == NULL)
                {
					q = &thisTree->nodes[taxonIndex];
					setAnc(q, p);
					setLft(p, q);
					p = q;
                }
				else if (getRht(p) == NULL)
                {
					q = &thisTree->nodes[taxonIndex];
					setAnc(q, p);
					setRht(p, q);
					p = q;
                }
				else if (getAnc(p) == NULL)
                {
					q = &thisTree->nodes[taxonIndex];
					setLft(q, p);
					setAnc(p, q);
					setFlag(q, true );
					p = q;
					thisTree->root = p;
                }
				else
                {
					cerr << "Problem finding an open node!" << endl;
					exit(1);
                }
				setIsLeaf(p, true );
				setName(p, (*s).c_str() );
            }
        }
    }
    
	if (thisTree->isRooted == true)
    {
		p = &thisTree->nodes[nextIntNode];
		q = thisTree->root;
		setLft(p, q );
		setAnc(q, p );
		thisTree->root = p;
    }
	/* initialize the post-order tree traversal sequence */
	getDownPassSequence(thisTree);
	for (int n=0; n<thisTree->numNodes; n++)
    {
		Node *p = getDownPassNode( thisTree, n );
		if (getAnc(p) != NULL)
			setV(p, thisTree->ranPtr->exponentialRv(10.0) );
		else
			setV(p, 0.0 );
    }
    
}

/* get the post-order tree traversal sequence */
void getDownPassSequence(Tree * thisTree) {
    
	int i = 0;
	passDn(thisTree, thisTree->root, &i);
    
}

Node   *getRoot(Tree * thisTree) { return thisTree->root; }
Node   *getDownPassNode(Tree * thisTree, int i) { return thisTree->downPassSequence[i]; }


void completeNode(Tree * thisTree, Node *p)
{
    if(getAnc(p) != NULL)
    {
        getAnc(p)->childrenLeft--;
        if(getAnc(p)->childrenLeft == 0)
        {
            enqueue(thisTree, getAnc(p));
        }
    }
}

void enqueue(Tree * thisTree, Node *p)
{
        thisTree->nextQueue[thisTree->queueEnd] = p;
        thisTree->queueEnd++;
}
Node *dequeue(Tree * thisTree)
{
    {
        if(thisTree->queueEnd != thisTree->queueStart)
        {
            return thisTree->nextQueue[thisTree->queueStart++];
        }
    }
    return NULL;
}
    

/* recursive function filling in the post-order traversal */
void passDn(Tree * thisTree, Node *p, int *x) {
    
	if (p != NULL)
    {
        if(getLft(p) != NULL) p->childrenLeft++;
        if(getRht(p) != NULL) p->childrenLeft++;
        if(p->childrenLeft == 0) enqueue(thisTree, p);
		passDn(thisTree, getLft(p), x);
		passDn(thisTree, getRht(p), x);
		thisTree->downPassSequence[(*x)++] = p;
        //fprintf(stderr,"Adding node %d with anc %d\n", *x, getAnc(p)!=NULL);
    }
    
}

/* print out the tree in an indented format */
void print(Tree * thisTree) {
    
	showNodes(thisTree, thisTree->root, 3);
    
}

/* show the nodes of the tree in an indented format */
void showNodes(Tree * thisTree, Node *p, int indent) {
    
	if (p != NULL)
    {
		for (int i=0; i<indent; i++)
			cout << " ";
		cout << dex(p) << " (" << dex(getLft(p)) << ", " << dex(getRht(p)) << ", " << dex(getAnc(p)) << ") ";
		if ( p != thisTree->root )
			cout << fixed << setprecision(6) << getV(p);
		if (getIsLeaf(p) == true)
			cout << " (" << getName(p) << ")";
		if (p == thisTree->root)
			cout << " <- Root" << endl;
		else
			cout << endl;
		showNodes (thisTree, getLft(p),  indent + 2);
		showNodes (thisTree, getRht(p), indent + 2);
    }
    
}

/* a small function that returns the index of the node (or -1 if the node is NULL) */
int dex(Node *p) {
    
	return ( p == NULL ? -1 : getIndex(p) );
    
}
